package crypto

import (
	"fmt"
	"sort"
)

// Polynomial represents a polynomial over a field
type Polynomial struct {
	Coeffs []ZR
	Field  string // "ZR" for scalar field, "G1" for group elements
}

// NewPolynomial creates a new polynomial with given coefficients
func NewPolynomial(coeffs []ZR) *Polynomial {
	// Remove trailing zeros
	trimmed := make([]ZR, 0, len(coeffs))
	for i := len(coeffs) - 1; i >= 0; i-- {
		if !coeffs[i].IsZero() {
			trimmed = coeffs[:i+1]
			break
		}
	}
	if len(trimmed) == 0 {
		trimmed = []ZR{{}}
	}
	
	return &Polynomial{
		Coeffs: trimmed,
		Field:  "ZR",
	}
}

// Evaluate evaluates the polynomial at point x
func (p *Polynomial) Evaluate(x ZR) ZR {
	var result ZR
	result.SetZero()
	
	var xPower ZR
	xPower.SetOne()
	
	for _, coeff := range p.Coeffs {
		var term ZR
		term.Mul(&coeff, &xPower)
		result.Add(&result, &term)
		
		var temp ZR
		temp.Mul(&xPower, &x)
		xPower = temp
	}
	
	return result
}

// Degree returns the degree of the polynomial
func (p *Polynomial) Degree() int {
	return len(p.Coeffs) - 1
}

// IsZero checks if the polynomial is zero
func (p *Polynomial) IsZero() bool {
	return len(p.Coeffs) == 0 || (len(p.Coeffs) == 1 && p.Coeffs[0].IsZero())
}

// Add adds two polynomials
func (p *Polynomial) Add(other *Polynomial) *Polynomial {
	maxLen := len(p.Coeffs)
	if len(other.Coeffs) > maxLen {
		maxLen = len(other.Coeffs)
	}
	
	result := make([]ZR, maxLen)
	
	for i := 0; i < maxLen; i++ {
		var sum ZR
		sum.SetZero()
		
		if i < len(p.Coeffs) {
			sum.Add(&sum, &p.Coeffs[i])
		}
		if i < len(other.Coeffs) {
			sum.Add(&sum, &other.Coeffs[i])
		}
		
		result[i] = sum
	}
	
	return NewPolynomial(result)
}

// Mul multiplies two polynomials
func (p *Polynomial) Mul(other *Polynomial) *Polynomial {
	if p.IsZero() || other.IsZero() {
		return NewPolynomial([]ZR{{}})
	}
	
	resultLen := p.Degree() + other.Degree() + 1
	result := make([]ZR, resultLen)
	
	for i := 0; i < resultLen; i++ {
		result[i].SetZero()
	}
	
	for i, coeffA := range p.Coeffs {
		for j, coeffB := range other.Coeffs {
			var product ZR
			product.Mul(&coeffA, &coeffB)
			result[i+j].Add(&result[i+j], &product)
		}
	}
	
	return NewPolynomial(result)
}

// LagrangeInterpolation performs Lagrange interpolation
func LagrangeInterpolation(points [][2]ZR, x ZR) (ZR, error) {
	if len(points) == 0 {
		return ZR{}, fmt.Errorf("no points provided")
	}
	
	var result ZR
	result.SetZero()
	
	for i, point := range points {
		xi, yi := point[0], point[1]
		
		// Compute Lagrange basis polynomial
		var numerator ZR
		numerator.SetOne()
		var denominator ZR
		denominator.SetOne()
		
		for j, otherPoint := range points {
			if i == j {
				continue
			}
			xj := otherPoint[0]
			
			// numerator *= (x - xj)
			var term ZR
			term.Sub(&x, &xj)
			numerator.Mul(&numerator, &term)
			
			// denominator *= (xi - xj)
			var denomTerm ZR
			denomTerm.Sub(&xi, &xj)
			denominator.Mul(&denominator, &denomTerm)
		}
		
		// Compute Lagrange coefficient: numerator / denominator
		var lagrangeCoeff ZR
		lagrangeCoeff.Div(&numerator, &denominator)
		
		// Add yi * lagrangeCoeff to result
		var term ZR
		term.Mul(&yi, &lagrangeCoeff)
		result.Add(&result, &term)
	}
	
	return result, nil
}

// InterpolateG1AtX performs Lagrange interpolation for G1 elements
func InterpolateG1AtX(coords [][2]interface{}, x int64, identity G1) (G1, error) {
	if len(coords) == 0 {
		return identity, fmt.Errorf("no coordinates provided")
	}
	
	// Convert coordinates to proper format
	points := make([][2]int64, len(coords))
	for i, coord := range coords {
		xi, ok1 := coord[0].(int64)
		_, ok2 := coord[1].(G1)
		if !ok1 || !ok2 {
			return identity, fmt.Errorf("invalid coordinate format at index %d", i)
		}
		points[i] = [2]int64{xi, xi} // Simplified for now
	}
	
	// Sort points by x coordinate
	sort.Slice(points, func(i, j int) bool {
		return points[i][0] < points[j][0]
	})
	
	// Use first f+1 points for interpolation
	order := len(points)
	if order > len(coords) {
		order = len(coords)
	}
	
	var result G1
	result = identity
	
	// Simplified interpolation - in practice this would involve proper G1 operations
	for i := 0; i < order; i++ {
		// This is a placeholder - actual implementation would involve
		// proper Lagrange interpolation with G1 elements
		var temp G1
		temp.SetInfinity()
		result = AddG1(&result, &temp)
	}
	
	return result, nil
}

// PolynomialInterpolateAtX performs Lagrange interpolation at point x (polynomial-specific version)
func PolynomialInterpolateAtX(coords [][2]interface{}, x int64) (ZR, error) {
	if len(coords) == 0 {
		return ZR{}, fmt.Errorf("no coordinates provided")
	}
	
	// Convert coordinates to proper format
	points := make([][2]ZR, len(coords))
	for i, coord := range coords {
		xi, ok1 := coord[0].(int64)
		yi, ok2 := coord[1].(ZR)
		if !ok1 || !ok2 {
			return ZR{}, fmt.Errorf("invalid coordinate format at index %d", i)
		}
		points[i] = [2]ZR{ZRFromInt(xi), yi}
	}
	
	// Sort points by x coordinate
	sort.Slice(points, func(i, j int) bool {
		return points[i][0].Element.Cmp(&points[j][0].Element) < 0
	})
	
	// Use first f+1 points for interpolation
	order := len(points)
	if order > len(coords) {
		order = len(coords)
	}
	
	// Perform Lagrange interpolation
	xZR := ZRFromInt(x)
	return LagrangeInterpolation(points[:order], xZR)
}

// CreateRandomPolynomial creates a random polynomial of given degree
func CreateRandomPolynomial(degree int, constantTerm ZR) []ZR {
	coeffs := make([]ZR, degree+1)
	coeffs[0] = constantTerm
	
	for i := 1; i <= degree; i++ {
		coeffs[i].SetRandom()
	}
	
	return coeffs
}

// EvaluatePolynomial evaluates a polynomial at point x (polynomial-specific version)
func EvaluatePolynomial(coeffs []ZR, x int64) ZR {
	var result ZR
	result.SetZero()
	
	xZR := ZRFromInt(x)
	var xPower ZR
	xPower.SetOne()
	
	for i := len(coeffs) - 1; i >= 0; i-- {
		var term ZR
		term.Mul(&coeffs[i], &xPower)
		result.Add(&result, &term)
		
		if i > 0 {
			var temp ZR
			temp.Mul(&xPower, &xZR)
			xPower = temp
		}
	}
	
	return result
}

// LagrangeCoefficient computes the Lagrange coefficient for interpolation
func LagrangeCoefficient(points []int64, j int, x int64) (ZR, error) {
	if j < 0 || j >= len(points) {
		return ZR{}, fmt.Errorf("index j out of range")
	}
	
	var numerator ZR
	numerator.SetOne()
	var denominator ZR
	denominator.SetOne()
	
	xZR := ZRFromInt(x)
	xjZR := ZRFromInt(points[j])
	
	for i, xi := range points {
		if i == j {
			continue
		}
		
		xiZR := ZRFromInt(xi)
		
		// numerator *= (x - xi)
		var term ZR
		term.Sub(&xZR, &xiZR)
		numerator.Mul(&numerator, &term)
		
		// denominator *= (xj - xi)
		var denomTerm ZR
		denomTerm.Sub(&xjZR, &xiZR)
		denominator.Mul(&denominator, &denomTerm)
	}
	
	// Return numerator / denominator
	var result ZR
	result.Div(&numerator, &denominator)
	return result, nil
}

// LagrangeCoefficientAtZero computes Lagrange coefficient for interpolation at 0
func LagrangeCoefficientAtZero(points []int64, j int) (ZR, error) {
	return LagrangeCoefficient(points, j, 0)
}

// InterpolateAtZero performs Lagrange interpolation at point 0
func InterpolateAtZero(points [][2]ZR) (ZR, error) {
	if len(points) == 0 {
		return ZR{}, fmt.Errorf("no points provided")
	}
	
	var result ZR
	result.SetZero()
	
	for i, point := range points {
		xi, yi := point[0], point[1]
		
		// Compute Lagrange coefficient at 0
		var numerator ZR
		numerator.SetOne()
		var denominator ZR
		denominator.SetOne()
		
		for j, otherPoint := range points {
			if i == j {
				continue
			}
			xj := otherPoint[0]
			
			// numerator *= (0 - xj) = -xj
			var negXj ZR
			negXj.Neg(&xj)
			numerator.Mul(&numerator, &negXj)
			
			// denominator *= (xi - xj)
			var denomTerm ZR
			denomTerm.Sub(&xi, &xj)
			denominator.Mul(&denominator, &denomTerm)
		}
		
		// Compute Lagrange coefficient: numerator / denominator
		var lagrangeCoeff ZR
		lagrangeCoeff.Div(&numerator, &denominator)
		
		// Add yi * lagrangeCoeff to result
		var term ZR
		term.Mul(&yi, &lagrangeCoeff)
		result.Add(&result, &term)
	}
	
	return result, nil
}
